﻿using RayDen.Library.Components.PhotonMapping;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;

namespace RayDen.RayEngine.Engines.PhotonMapping
{

    public class PhotonMappingPathSampler : PathSamplerBase, IContextSampler<PhotonMappingContext>
    {

        struct ShadowRayInfo
        {
            public RayData ShadowRay;
            public int ShadowRayIndex;
            public RgbSpectrum Throughput;
            public float Pdf;
        }

        struct FinalGatherRayInfo
        {
            public RayData Ray;
            public int RayIndex;
        }


        internal PhotonMappingContext context;
        internal RgbSpectrum Throughput, DirectBsdf;
        private bool shadowRayMode;
        private int tracedShadowRays;
        private Point hitPoint;
        private Normal shadingNormal;
        private ShadowRayInfo[] shadowRays;
        private FinalGatherRayInfo[] fgRays;
        private RayEngineScene scene;
        private int depth;
        private bool specularBounce;


        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Sample = pathIntegrator.Sampler.GetSample(null);
            this.Radiance = new RgbSpectrum(0.000f);
            this.DirectBsdf = RgbSpectrum.ZeroSpectrum();
            this.Throughput = RgbSpectrum.UnitSpectrum();
            this.shadowRayMode = false;
            this.tracedShadowRays = 0;
            this.specularBounce = true;
            this.depth = 0;
            this.specularBounce = false;
            if (PhotonMappingContext.Mode == PhotonMappingMode.FinalGathering && this.fgRays == null)
            {
                fgRays = new FinalGatherRayInfo[context.FinalGatheringRays];
            }
            if (this.shadowRays == null)
                this.shadowRays = new ShadowRayInfo[scene.ShadowRaysPerSample];
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out this.PathRay);
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            if (shadowRayMode)
            {
                if (PhotonMappingContext.Mode == PhotonMappingMode.FinalGathering)
                {
                    if (rayBuffer.LeftSpace() > fgRays.Length)
                    {
                        for (int index = 0; index < shadowRays.Length; index++)
                        {
                            fgRays[index].RayIndex = rayBuffer.AddRay(ref fgRays[index].Ray);
                        }
                    }
                    else
                    {
                        return false;
                    }
                }
                if (rayBuffer.LeftSpace() > tracedShadowRays)
                {
                    for (int index = 0; index < tracedShadowRays; index++)
                    {
                        shadowRays[index].ShadowRayIndex = rayBuffer.AddRay(ref shadowRays[index].ShadowRay);
                    }
                    return true;
                }
                return false;
            }
            if (rayBuffer.LeftSpace() < 1)
                return false;
            RayIndex = rayBuffer.AddRay(ref PathRay);
            return true;
        }

        SurfaceIntersectionData hitInfo = null;
        public override void Advance(Core.Types.RayBuffer rayBuffer, Core.Types.SampleBuffer consumer)
        {
#if VERBOSE

            try
            {
#endif
            Vector wo = -PathRay.Dir;
            if (shadowRayMode)
            {
                for (int i = 0; i < tracedShadowRays; i++)
                {
                    RgbSpectrum att;

                    bool continueTrace;
                    if (ShadowRayTest(ref rayBuffer.rayHits[shadowRays[i].ShadowRayIndex], ref shadowRays[i].ShadowRay, out att, out continueTrace))
                    {
                        Radiance += att * shadowRays[i].Throughput / shadowRays[i].Pdf;
                    }
                    switch (PhotonMappingContext.Mode)
                    {
                        case PhotonMappingMode.SeparateEstimation:
                                  this.Radiance += Throughput*EvaluateIndirectRadiance(context.globalPhotonMap, context.indirectRadius, context.indirectPhotons,
                                                                null, ref hitPoint, ref wo,
                                                                ref shadingNormal);

                            /*
                                  this.Radiance += Throughput * EvaluateIndirectRadiance(context.causticPhotonMap, context.causticRadius, context.causticPhotons, null, ref hitPoint, ref wo,ref shadingNormal);
                             */ 

                            break;
                        case PhotonMappingMode.FinalGathering:
                            var indirect = RgbSpectrum.ZeroSpectrum();
                            for (int index = 0; index < fgRays.Length; index++)
                            {
                                var fgHit = rayBuffer.rayHits[fgRays[index].RayIndex];
                                if (fgHit.Miss())
                                    continue;
                                var fgPoint = fgRays[index].Ray.Point(fgHit.Distance);
                                hitInfo = SurfaceSampler.GetIntersection(ref  fgRays[index].Ray, ref fgHit);
                                if (hitInfo.MMaterial.IsDiffuse())
                                {
                                    
                                    indirect += EvaluateIndirectRadiance(context.globalPhotonMap, context.finalGatheringRadius,
                                                                 context.finalGatheringPhotons,
                                                                 hitInfo.MMaterial, ref fgPoint, ref wo,
                                                                 ref shadingNormal);
                                    /*
                                    RgbSpectrum fs;
                                    hitInfo.MMaterial.f(
                                        ref fgRays[index].Ray.Dir,
                                        ref wo, ref hitInfo.ShadingNormal, out fs, BrdfType.Diffuse);
                                    indirect += context.globalPhotonMap.IrradianceEstimate(ref fgPoint,
                                                                                           ref hitInfo.ShadingNormal,
                                                                                           context.finalGatheringRadius,
                                                                                           context.finalGatheringPhotons);
                                     */ 
                                }
                            }

                            //indirect +=context.causticPhotonMap.IrradianceEstimate(ref hitPoint, ref shadingNormal, .2f, 10);
                            float weight = 1f/ (context.FinalGatheringRays);
                            Radiance += indirect * weight*Throughput;
                            break;
                        case PhotonMappingMode.DirectOnly:
                            break;
                    }
                    //Use Maps Here
                }
                shadowRayMode = false;
                tracedShadowRays = 0;
                Splat(consumer);
                return;
            }

            var rayHit = rayBuffer.rayHits[RayIndex];
            bool missed = rayHit.Index == 0xffffffffu;
            if (missed || depth > scene.MaxPathDepth)
            {
                if (depth < 2 || specularBounce)
                {
                    this.Radiance += Throughput * this.LightSampler.SampleEnvironment(-PathRay.Dir);
                }
                Splat(consumer);
                return;
            }



            depth++;

            if (hitInfo == null)
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
            else
            {
                SurfaceSampler.GetIntersection(ref PathRay, ref rayHit, ref hitInfo);
            }
            hitPoint = PathRay.Point(rayHit.Distance);
            shadingNormal = hitInfo.ShadingNormal;
            var bsdf = hitInfo.MMaterial;


            if (context.mapPreviewMode)
            {

                //this.Radiance  += context.causticPhotonMap.IrradianceEstimate(ref hitPoint, ref shadingNormal, .2f, 30);
                //this.Radiance = context.globalPhotonMap.IrradianceEstimate(ref hitPoint, ref shadingNormal, context.indirectRadius, context.indirectPhotons);
                
                /*
                this.Radiance = EvaluateIndirectRadiance(context.globalPhotonMap, context.finalGatheringRadius, context.finalGatheringPhotons,
                                                                hitInfo.MMaterial, ref hitPoint, ref wo,
                                                                ref shadingNormal);
                */
                
                this.Radiance = EvaluateIndirectRadiance(context.globalPhotonMap, context.indirectRadius, context.indirectPhotons,
                                                                hitInfo.MMaterial, ref hitPoint, ref wo,
                                                                ref shadingNormal);
                
          /*
                this.Radiance = EvaluateIndirectRadiance(context.causticPhotonMap, context.causticRadius, context.causticPhotons,
                                                               hitInfo.MMaterial, ref hitPoint, ref wo,
                                                                ref shadingNormal);*/
                this.Splat(consumer);
                return;
            }

            if (hitInfo.IsLight)
            {
                //if (specularBounce)
                {
                    var lt = scene.GetLightByIndex((int)rayHit.Index);
                    if (lt != null)
                    {
                        var mwi = -PathRay.Dir;
                        var le = (RgbSpectrum)(RgbSpectrumInfo)(lt.Le(ref mwi));
                        Radiance += Throughput * le;
                    }
                }
                Splat(consumer);
                return;
            }



            if (bsdf.IsDiffuse())
            {
                float lightStrategyPdf = LightSampler.StrategyPdf;
                //scene.ShadowRaysPerSample/(float)scene.Lights.Length;
                RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                LightSampler.EvaluateShadow(ref hitPoint, ref hitInfo.Normal, Sample.GetLazyValue(),
                                            Sample.GetLazyValue(), Sample.GetLazyValue(), ref ls);
                foreach (var lightSample in ls)
                {
                    if (lightSample.Pdf <= 0f)
                        continue;
                    shadowRays[tracedShadowRays].Throughput = (RgbSpectrum)(RgbSpectrumInfo)(lightSample.Spectrum);
                    shadowRays[tracedShadowRays].Pdf = lightSample.Pdf;
                    shadowRays[tracedShadowRays].ShadowRay = lightSample.LightRay;
                    Vector lwi = shadowRays[tracedShadowRays].ShadowRay.Dir;
                    RgbSpectrum fs;
                    hitInfo.MMaterial.f(
                        ref shadowRays[tracedShadowRays].ShadowRay.Dir,
                        ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs, types: BrdfType.Diffuse);
                    shadowRays[tracedShadowRays].Throughput *= lightTroughtput *
                                                           Vector.AbsDot(ref hitInfo.Normal, ref lwi) *
                                                           fs;

                    if (!shadowRays[tracedShadowRays].Throughput.IsBlack())
                    {
                        shadowRays[tracedShadowRays].Pdf /= lightStrategyPdf;
                        tracedShadowRays++;
                    }
                }

                if (PhotonMappingContext.Mode == PhotonMappingMode.FinalGathering)
                {
                    var onb = new ONB();
                    onb.SetFromZ(ref hitInfo.Normal);

                    /*
                    Vector v1, v2;
                    Normal n = N;
                    Vector.CoordinateSystem(ref n, out v1, out v2);

                    dir = new Vector(
                        v1.x * dir.x + v2.x * dir.y + n.x * dir.z,
                        v1.y * dir.x + v2.y * dir.y + n.y * dir.z,
                        v1.z * dir.x + v2.z * dir.y + n.z * dir.z);
                     */
                    for (int i = 0; i < context.FinalGatheringRays; i++)
                    {
                        var dir = MC.CosineSampleHemisphere(Sample.GetLazyValue(), Sample.GetLazyValue());
                        onb.ToWorld(ref dir);
                        fgRays[i] = new FinalGatherRayInfo() { Ray = new RayData(ref hitPoint, ref dir) };

                    }
                }

                if (tracedShadowRays > 0)
                {
                    shadowRayMode = true;
                    return;
                }
            }



            float fPdf;
            Vector wi;

            RgbSpectrum f = hitInfo.MMaterial.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                                      Sample.GetLazyValue(), Sample.GetLazyValue(),
                                                      Sample.GetLazyValue(), ref hitInfo.TextureData,
                                                      out fPdf, out specularBounce) * hitInfo.Color;

            if (f.IsBlack() || fPdf <= 0f && !shadowRayMode)
            {
                Splat(consumer);
                return;
            }
            if (bsdf.IsDiffuse())
                DirectBsdf = f / fPdf;

            Throughput *= f / fPdf;
            PathRay = new RayData(ref hitPoint, ref wi);

#if VERBOSE
            }
            catch (System.Exception ex)
            {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }
        const float k = 1.0f;

        private const float a = 0.918f;
        private const float b = 1.953f;

        private static float GaussKernel( float distance, float radius)
        {
            //return 1f; //Disk
            //return 1f - (distance/(radius*k)); //Cone


            return a*(1.0f - ((1.0f - MathLab.Exp((-b*distance*distance) / (2.0f* radius*radius)))/(1.0f - MathLab.Exp(-b))));
        }

        private static float ConeKernel(float distance, float radius)
        {
            return 1f - (distance/(radius*k)); //Cone
        }


        

        public enum PhotonMapFilteringMode
        {
            None,
            Disk,
            Cone,
            Gauss
        }

        private RgbSpectrum EstimateRadiance(Photon[] photons, float radius, ref Point hitPt, ref Vector wo, ref Normal ns, RgbSpectrum f, ISurfaceMaterial bsdf = null, PhotonMapFilteringMode mode = PhotonMapFilteringMode.Cone)
        {
            float totalWeight = 1f;
            var radiance = new RgbSpectrum(0f);

            switch (mode)
            {
                case PhotonMapFilteringMode.Disk:
                    totalWeight = (MathLab.INVPI)/(radius*radius);
                    break;
                case PhotonMapFilteringMode.Cone:
                    totalWeight = (1f / ((1.0f - (2.0f / 3f * k)) * radius * radius * MathLab.M_PI));
                    break;
            }

            if (photons.Length > 0)
            {
                for (int index = 0; index < photons.Length; index++)
                {
                    var photon = photons[index];
                    Vector pdir = KdTreePhotonMap.PhotonDirection(ref photon);
                    var dot = Vector.Dot(ref pdir, ref ns);

                    if (dot > 0.0f)
                    {
                        float ws = 1.0f;
                        switch (mode)
                        {
                            case PhotonMapFilteringMode.Gauss:
                                ws = GaussKernel((photon.Position - hitPt.ToVec()).Length, radius);
                                break;
                            case PhotonMapFilteringMode.Cone:
                                ws = ConeKernel((photon.Position - hitPt.ToVec()).Length, radius);
                                break;
                        }
                        RgbSpectrum fs = f;
                        if (bsdf != null)
                        {
                            bsdf.f(ref wo,ref pdir, ref ns, ref Throughput, out fs);
                        }
                        radiance += photon.Power * f * ws;
                    }
                }
                return radiance*totalWeight;
                //radiance*= (MathLab.INVPI)/(radius*radius); //Disk

                //radiance *= (1f / ((1.0f - (2.0f/3f*k))*radius*radius*MathLab.M_PI)); // Cone
                //None modification for gauss
            }
            return RgbSpectrum.ZeroSpectrum();
        }

      


        private RgbSpectrum EvaluateIndirectRadiance(IPhotonMap map , float radius, int maxPhotons ,ISurfaceMaterial bsdf, ref Point hitPt,ref Vector wo, ref Normal ns)
        {
            Photon[] photons;
            var radiance = RgbSpectrum.ZeroSpectrum();
            map.QueryMap(ref hitPt, radius, maxPhotons, out photons);
            radiance += EstimateRadiance(photons, radius, ref hitPt, ref wo, ref ns, RgbSpectrum.UnitSpectrum(), bsdf,
                                         PhotonMapFilteringMode.Cone);
            return radiance;
        }

        public void SetContext(PhotonMappingContext context)
        {
            this.context = context;
        }
    }
}
